import torch
from math import inf
from torch import nn

QK_val = torch.ones((2, 5, 5))

# ================ 手写掩码 ======================
# atten_mask = torch.tensor([
#     [0, float(-inf), float(-inf), float(-inf), float(-inf)],
#     [0, 0, float(-inf), float(-inf), float(-inf)],
#     [0, 0, 0, float(-inf), float(-inf)],
#     [0, 0, 0, 0, float(-inf)],
#     [0, 0, 0, 0, 0],
# ]).float()
#
# print(torch.softmax((QK_val + atten_mask), dim=-1))

# ================ 使用torch.triu设计斜对角线矩阵 =================
# atten_mask = torch.ones(QK_val.shape)
# mask_triu = torch.triu(atten_mask, 1)
# mask_triu[mask_triu == 1] = float(-inf)
# atten_mask = mask_triu
# print(torch.softmax((QK_val + atten_mask), dim=-1))


# ============= nn.Transformer.generate_square_subsequent_mask 直接生成掩码矩阵 =========
atten_mask = nn.Transformer.generate_square_subsequent_mask(5)
print(torch.softmax((QK_val + atten_mask), dim=-1))
